#ifndef TEST_CNN_MPI_BGV_MNIST
#define TEST_CNN_MPI_BGV_MNIST

#include "CNN/mpi_mnist.hpp"
#include "crypto_helper.hpp"
#include "keys.h"

#define T_NUM 4
#define T_BATCHSIZE 64
#define conv_size 49

static constexpr double SCALE = 1024.0;

static constexpr unsigned long MATRIX_DIM_M = 10;
static constexpr unsigned long MATRIX_DIM_N = 256;
static constexpr unsigned long MATRIX_DIM_K = 64;
// static constexpr unsigned long MATRIX_DIM_M = 64;
// static constexpr unsigned long MATRIX_DIM_N = 64;
// static constexpr unsigned long MATRIX_DIM_K = 64;
static constexpr unsigned long HYPERCUBE_DIM_X =
    MAX(MATRIX_DIM_M, MATRIX_DIM_K);
static constexpr unsigned long HYPERCUBE_DIM_Y =
    MAX(MATRIX_DIM_K, MATRIX_DIM_N);

// std::vector<helib::PubKey* > global_pk;

// double bcast_pubkey(int root, int rank, MPI_Comm comm) {
//   std::string pk_str;
//   std::stringstream pk_ss;
//   auto bgv_client = static_cast<mpi_bgv_hmnist_client *>(
//       const_cast<mpi_h_mnist_client<helib::BGV> *>(client_));

//   if (rank == root) {
//     global_pk->writeTo(pk_ss);
//     pk_str = pk_ss.str();
//   }
//   MPI_Bcast_String(pk_str, root, comm);
//   if (rank != root) {
//     pk_ss << pk_str;
//   }
//   return pk_str.size() / 1024.0 / 1024.0;
// }

static cube_signature MNIST_HYPERCUBE({HYPERCUBE_DIM_X, HYPERCUBE_DIM_Y});

class mpi_bgv_hmnist_client : public mpi_h_mnist_client<helib::BGV> {
private:
  mpi_bgv_hmnist_client() = delete;

  helib::Ctxt init_hx(const torch::Tensor &a, int asz0, int asz1,
                      int t_index) const;

  helib::Ctxt init_hconv(const torch::Tensor &a, int asz0, int asz1,
                         int t_index) const;

  helib::Ctxt init_hfc1(const torch::Tensor &a, int asz0, int t_index) const;

  helib::Ctxt init_hfc2(const torch::Tensor &a, int t_index) const;

public:
  mpi_bgv_hmnist_client(const params<helib::BGV> &params, int crt_num);
  ~mpi_bgv_hmnist_client() {
    for (auto hmmcc : hmmccs_) {
      delete hmmcc;
    }
    for (auto hmme : hmmes_) {
      delete hmme;
    }
  }

  // double bcast_pubkey(int root, int rank, MPI_Comm comm, helib::SecKey *sk);

  int get_crt_num() const { return crt_num_; }

  const std::vector<helib::Context *> &get_context() const { return contexts_; }

  const std::vector<hypercube_hmmcc *> &get_hmmcc() const { return hmmccs_; }

  const std::vector<hypercube_hmme *> &get_hmme() const { return hmmes_; }
  // Generate encrypted model
  void init_hcnn(encrypted_model<helib::BGV> &ctxt_model,
                 const torch::nn::Conv2d &conv1, const torch::nn::Linear &fc1,
                 const torch::nn::Linear &fc2) const override;

  // Generate encrypted image
  void init_input(encrypted_input<helib::BGV> &ctxt_input,
                  const torch::Tensor &filter,
                  const torch::Tensor &x) const override;

  void recover_result(at::Tensor &out, const std::vector<std::string> &in,
                      MPI_Comm comm, int root) const override;

private:
  ENC key = ENC::PUBKEY;

  int crt_num_;
  NTL::ZZ T_;
  std::vector<NTL::ZZ> ti_;

  std::vector<hypercube_hmmcc *> hmmccs_;
  std::vector<helib::Context *> contexts_;
  std::vector<hypercube_hmme *> hmmes_;
};

class mpi_bgv_hmnist_server : public mpi_h_mnist_server<helib::BGV> {
private:
  mpi_bgv_hmnist_server() = delete;

public:
  std::shared_ptr<helib::Context>
  build_context(const params<helib::BGV> &params, long m);

  mpi_bgv_hmnist_server(const hmm_status<shmm_engine> &status, MPI_Comm comm,
                        int crt_num, int conv_size0,
                        const params<helib::BGV> &params);

  mpi_bgv_hmnist_server(const hmm_status<shmm_engine> &status, MPI_Comm comm,
                        int crt_num, int conv_size0,
                        const params<helib::BGV> &params,
                        const std::vector<hypercube_hmmcc *> &hmmcc);
  ~mpi_bgv_hmnist_server() {
    // for (auto hmme : hmmes_) {
    //   delete hmme;
    // }
    // delete contexts_;
    delete hmmes_;
    // delete hmmcss_;
    // delete pk_;
    for (auto pk : pk_) {
      delete pk;
    }
  }
  const std::vector<helib::Context *> &get_context() const { return contexts_; }
  // const helib::Context *get_context() const { return contexts_; }

  int get_crt_num() const { return crt_num_; }

  const hypercube_hmme *get_hmme() const { return hmmes_; }
  // const std::vector<helib::PubKey *> &get_pk() const { return pk_; }

  const std::vector<hypercube_hmmcs *> &get_hmmcs() const { return hmmcss_; }
  // const hypercube_hmmcs *get_hmmcs() const { return hmmcss_; }

  const helib::PubKey &get_pk(int i) const { return *pk_[i]; }

  void forward_test(std::vector<std::string> &out,
                    std::vector<std::string> &buf_input,
                    const std::vector<helib::Ctxt> &hconv,
                    const helib::Ctxt &hfc1,
                    const helib::Ctxt &hfc2) const override;

  void report_metrics(std::ostream &ss = std::cout) const override;

  void send_input_p2p(const encrypted_input<helib::BGV> &ctxt_in, MPI_Comm comm,
                      std::vector<std::string> &buf, int root) const override;

  void
  send_honv_p2p(const std::vector<std::vector<std::vector<helib::Ctxt>>> &ctxt,
                MPI_Comm comm, std::vector<std::string> &buf, int root,
                int rank, int crt_num) const override;
  // void back_honv_p2p(const std::vector<std::vector<std::string>> &ctxt,
  //                    MPI_Comm comm,
  //                    std::vector<std::vector<std::vector<std::string>>> &buf,
  //                    int root, int rank, int crt_num) const override;
  void send_hfc1_p2p(std::vector<std::vector<helib::Ctxt>> &ctxt, MPI_Comm comm,
                     std::string &buf, int root, int rank,
                     int crt_num) const override;
  // void back_hfc1_p2p(const std::vector<std::string> &ctxt, MPI_Comm comm,
  //                    std::vector<std::vector<std::string>> &buf, int root,
  //                    int rank, int crt_num) const override;

  void send_hfc2_p2p(const std::vector<helib::Ctxt> &ctxt, MPI_Comm comm,
                     std::string &buf, int root, int rank) const override;

  // void bcast_hconv_server(const std::vector<std::vector<std::string>> &in,
  //                         MPI_Comm comm,
  //                         std::vector<std::vector<std::string>> &buf, int
  //                         root, int rank, int crt_num) const override;
  // void
  // back_hconv_server(const std::vector<std::vector<std::string>> &ctxt,
  //                   MPI_Comm comm,
  //                   std::vector<std::vector<std::vector<std::string>>> &buf,
  //                   int root, int rank, int crt_num) const override;

  // void bcast_hfc1_server(const std::vector<std::string> &in, MPI_Comm comm,
  //                        std::vector<std::string> &buf, int root, int rank,
  //                        int crt_num) const override;
  // void
  // back_input_server(const std::vector<std::vector<std::string>> &ctxt,
  //                   MPI_Comm comm,
  //                   std::vector<std::vector<std::vector<std::string>>> &buf,
  //                   int root, int rank, int crt_num) const override;

  void hconv_server(helib::Ctxt &out, const std::vector<helib::Ctxt> &in,
                    const std::vector<helib::Ctxt> &hconv) const;
  void back_hconv(const std::vector<helib::Ctxt> &out,
                  std::vector<std::vector<helib::Ctxt>> &ctxt) const;
  void hsqr1_server(helib::Ctxt &ctxt) const;
  void hfc1_server(helib::Ctxt &ctxt, const helib::Ctxt &hfc1, MPI_Comm comm,
                   int rank, int root) const;
  void back_hfc1(const std::vector<helib::Ctxt> &out,
                 std::vector<std::vector<helib::Ctxt>> &ctxt,
                 MPI_Comm comm) const;
  void hsqr2_server(helib::Ctxt &ctxt) const;
  void hfc2_server(helib::Ctxt &ctxt, std::vector<std::string> &buf,
                   const helib::Ctxt &hfc2, MPI_Comm comm, int rank,
                   int root) const;

public:
  std::vector<helib::PubKey *> pk_;

private:
  std::vector<hypercube_hmmcs *> hmmcss_;

  // std::vector<metric> metrics_;
  std::vector<helib::Context *> contexts_;
  // int crt_num_;
  // std::vector<hypercube_hmme *> hmmes_;
  // hypercube_hmmcs *hmmcss_ = nullptr;

  std::vector<metric> metrics_;
  // helib::Context *contexts_ = nullptr;
  int crt_num_;
  hypercube_hmme *hmmes_ = nullptr;
};

#endif /* TEST_CNN_BGV_MNIST_HELPER */
